;;; PLT Port of Stephen Adams's Weight Balanced Trees

; This is a straightforward port of Stephen Adams portable implementation
; of weight balanced trees to the PLT module system.
;   /Jens Axel S?gaard,  16th nov 2003

; CHANGES
;  - 

;; -*-Scheme-*-
;;
;;  $Id: wttree.ss,v 1.1 1994/11/28 21:58:48 adams Exp adams $
;;
;;  Copyright (c) 1993-1994 Stephen Adams
;;
;;  References:
;;
;;    Stephen Adams, Implemeting Sets Efficiently in a Functional
;;       Language, CSTR 92-10, Department of Electronics and Computer
;;       Science, University of Southampton, 1992
;;
;;
;;  Copyright (c) 1993-94 Massachusetts Institute of Technology
;;
;;  This material was developed by the Scheme project at the Massachusetts
;;  Institute of Technology, Department of Electrical Engineering and
;;  Computer Science.  Permission to copy this software, to redistribute
;;  it, and to use it for any purpose is granted, subject to the following
;;  restrictions and understandings.
;;
;;  1. Any copy made of this software must include this copyright notice
;;  in full.
;;
;;  2. Users of this software agree to make their best efforts (a) to
;;  return to the MIT Scheme project any improvements or extensions that
;;  they make, so that these may be included in future releases; and (b)
;;  to inform MIT of noteworthy uses of this software.
;;
;;  3. All materials developed as a consequence of the use of this
;;  software shall duly acknowledge such use, in accordance with the usual
;;  standards of acknowledging credit in academic research.
;;
;;  4. MIT has made no warrantee or representation that the operation of
;;  this software will be error-free, and MIT is under no obligation to
;;  provide any services, by way of maintenance, update, or otherwise.
;;
;;  5. In conjunction with products arising from the use of this material,
;;  there shall be no use of the name of the Massachusetts Institute of
;;  Technology nor of any adaptation thereof in any advertising,
;;  promotional, or sales literature without prior written consent from
;;  MIT in each case.


(module set-weight-balanced-tree mzscheme
  (provide weight-balanced-tree-set@)

  (require (lib "unitsig.ss")
           (prefix list: (lib "list.ss"))
           (lib "set-signature.ss" "galore" "units")
           (lib "ordered.ss" "galore" "units"))

  (define weight-balanced-tree-set@
    (unit/sig set^
      (import ordered^)
      
      (define-struct node (k l r w))
      
      (define empty  'empty)
      (define (empty? x) (eq? x 'empty))
      
      (define (size node)
        (if (empty? node) 0  (node-w node)))
      
      (define (set . args) (list:foldl (lambda (x acc) (insert acc x)) empty args))
      (define set? node?)
      
      (define (singleton k) (make-node k empty empty 1))
      
      ;;  (define (with-n-node node receiver)
      ;;    (receiver (node/k node) (node/v node) (node/l node) (node/r node)))
      
      (define-syntax with-n-node
        (syntax-rules ()
          [(with-n-node node receiver) (receiver (node-k node) (node-l node) (node-r node))]))
      
      ;;
      ;;  Constructors for building node trees of various complexity
      ;;
      
      (define (n-join k l r)
        (make-node k l r (add1 (+ (size l) (size r)))))
      
      (define (single-l a.k x r)
        (with-n-node r
                     (lambda (b.k y z) (n-join b.k (n-join a.k x y) z))))
      
      (define (double-l a.k x r)
        (with-n-node r
                     (lambda (c.k r.l z)
                       (with-n-node r.l
                                    (lambda (b.k y1 y2)
                                      (n-join b.k 
                                              (n-join a.k x y1)
                                              (n-join c.k y2 z)))))))
      
      (define (single-r b.k l z)
        (with-n-node l
                     (lambda (a.k x y) (n-join a.k x (n-join b.k y z)))))
      
      (define (double-r c.k l z)
        (with-n-node l
                     (lambda (a.k x l.r)
                       (with-n-node l.r
                                    (lambda (b.k y1 y2)
                                      (n-join b.k
                                              (n-join a.k x y1)
                                              (n-join c.k y2 z)))))))
      
      ;; (define-integrable wt-tree-ratio 5)
      (define wt-tree-ratio 5)
      
      (define (t-join k l r)
        (define (simple-join) (n-join k l r))
        (let ((l.n  (size l))
              (r.n  (size r)))
          (cond ((< (+ l.n r.n) 2)   (simple-join))
                ((> r.n (* wt-tree-ratio l.n))
                 ;; right is too big
                 (let ((r.l.n  (size (node-l r)))
                       (r.r.n  (size (node-r r))))
                   (if (< r.l.n r.r.n)
                       (single-l k l r)
                       (double-l k l r))))
                ((> l.n (* wt-tree-ratio r.n))
                 ;; left is too big
                 (let ((l.l.n  (size (node-l l)))
                       (l.r.n  (size (node-r l))))
                   (if (< l.r.n l.l.n)
                       (single-r k l r)
                       (double-r k l r))))
                (else
                 (simple-join)))))
      ;;
      ;;  Node tree procedures that are independent of elm<
      ;;
      
      (define (node/min node)
        (cond  ((empty? node)          (error:empty 'min))
               ((empty? (node-l node)) node)
               (else                   (node/min (node-l node)))))
      
      (define (node/delmin node)
        (cond ((empty? node)           (error:empty 'delmin))
              ((empty? (node-l node))  (node-r node))
              (else   (t-join (node-k node)
                              (node/delmin (node-l node)) (node-r node)))))
      
      (define (node/concat2 node1 node2)
        (cond ((empty? node1)   node2)
              ((empty? node2)   node1)
              (else
               (let ((min-node (node/min node2)))
                 (t-join (node-k min-node)
                         node1 (node/delmin node2))))))
      
      (define (fold procedure base node)
        (if (empty? node)
            base
            (with-n-node node
                         (lambda (k l r)
                           (fold procedure (procedure k (fold procedure base r)) l)))))
      
      (define (node/for-each procedure node)
        (if (not (empty? node))
            (with-n-node node
                         (lambda (k l r)
                           (node/for-each procedure l)
                           (procedure k)
                           (node/for-each procedure r)))))
      
      (define (node/height node)
        (if (empty? node)
            0
            (add1 (max (node/height (node-l node))
                       (node/height (node-r node))))))
      
      (define (node/index node index)
        (define (loop node index)
          (let ((size.l  (size (node-l node))))
            (cond ((< index size.l)  (loop (node-l node) index))
                  ((> index size.l)  (loop (node-r node)
                                               (- index (add1 size.l))))
                  (else                  node))))
        (let ((bound  (size node)))
          (if (or (< index 0)
                  (>= index bound))
              (error "node/index: Index out of range." index)
              (loop node index))))
      
      (define (error:empty owner)
        (error "Operation requires non-empty tree:" owner))
      
      
      (define (elm> x y)  (elm< y x))
      
      (define (node/member node k)
        ;; Returns either the node or #f.
        ;; Loop takes D comparisons where D is the depth of the tree
        ;; rather than the traditional compare-low, compare-high which
        ;; takes on average 1.5(D-1) comparisons
        (define (loop this best)
          (cond ((empty? this)  best)
                ((elm< k (node-k this))   (loop (node-l this) best))
                (else (loop (node-r this) this))))
        (let ((best (loop node #f)))
          (cond ((not best)               #f)
                ((elm< (node-k best) k)  #f)
                (else                     best))))
      
      (define (get node k) (node-k (node/member node k)))
      (define (member? node k) (and (node/member node k) #t))
      
      (define (node/rank k node rank)
        (cond ((empty? node)             #f)
              ((elm< k (node-k node))  (node/rank k (node-l node) rank))
              ((elm> k (node-k node))  
               (node/rank k (node-r node)
                          (+ 1 (+ rank (size (node-l node))))))
              (else                     (+ rank (size (node-l node))))))
    
      (define (insert-all node ks)
        (list:foldl (lambda (k acc) (insert acc k)) node ks))
      
      (define (insert node k)
        (if (empty? node)
            (singleton k)
            (with-n-node node
                         (lambda (key l r)
                           (cond ((elm< k key)   (t-join key (insert l k) r))
                                 ((elm< key k)   (t-join key l (insert r k)))
                                 (else            (n-join key l r)))))))
      
      (define (remove node x)
        (if (empty? node)
            empty
            (with-n-node node
                         (lambda (key l r)
                           (cond ((elm< x key)   (t-join key (remove l x) r))
                                 ((elm< key x)   (t-join key l (remove r x)))
                                 (else            (node/concat2 l r)))))))
      
      (define (remove-all node xs)
        (list:foldl (lambda (x acc) (remove acc x)) node xs))
      
      (define (node/concat tree1 tree2)
        (cond ((empty? tree1)  tree2)
              ((empty? tree2)  tree1)
              (else
               (let ((min-node (node/min tree2)))
                 (node/concat3 (node-k min-node) tree1
                               (node/delmin tree2))))))
      
      (define (node/concat3 k l r)
        (cond ((empty? l)   (insert r k))
              ((empty? r)   (insert l k))
              (else
               (let ((n1  (size l))
                     (n2  (size r)))
                 (cond ((< (* wt-tree-ratio n1) n2)
                        (with-n-node r
                                     (lambda (k2 l2 r2)
                                       (t-join k2 (node/concat3 k l l2) r2))))
                       ((< (* wt-tree-ratio n2) n1)
                        (with-n-node l
                                     (lambda (k1 l1 r1)
                                       (t-join k1 l1 (node/concat3 k r1 r)))))
                       (else
                        (n-join k l r)))))))
      
      (define (node/split-lt node x)
        (cond ((empty? node)  empty)
              ((elm< x (node-k node))
               (node/split-lt (node-l node) x))
              ((elm< (node-k node) x)
               (node/concat3 (node-k node) (node-l node)
                             (node/split-lt (node-r node) x)))
              (else (node-l node))))
      
      (define (node/split-gt node x)
        (cond ((empty? node)  empty)
              ((elm< (node-k node) x)
               (node/split-gt (node-r node) x))
              ((elm< x (node-k node))
               (node/concat3 (node-k node) 
                             (node/split-gt (node-l node) x) (node-r node)))
              (else (node-r node))))
      
      (define (union tree1 tree2)
        (cond ((empty? tree1)  tree2)
              ((empty? tree2)  tree1)
              (else
               (with-n-node tree2
                            (lambda (ak l r)
                              (let ((l1  (node/split-lt tree1 ak))
                                    (r1  (node/split-gt tree1 ak)))
                                (node/concat3 ak (union l1 l) (union r1 r))))))))
      
      (define (difference tree1 tree2)
        (cond ((empty? tree1)   empty)
              ((empty? tree2)   tree1)
              (else
               (with-n-node tree2
                            (lambda (ak l r)
                              (let ((l1  (node/split-lt tree1 ak))
                                    (r1  (node/split-gt tree1 ak)))
                                (node/concat (difference l1 l)
                                             (difference r1 r))))))))
      
      (define (intersection tree1 tree2)
        (cond ((empty? tree1)   empty)
              ((empty? tree2)   empty)
              (else
               (with-n-node tree2
                            (lambda (ak l r)
                              (let ((l1  (node/split-lt tree1 ak))
                                    (r1  (node/split-gt tree1 ak)))
                                (if (member? tree1 ak)
                                    (node/concat3 ak (intersection l1 l)
                                                  (intersection r1 r))
                                    (node/concat (intersection l1 l)
                                                 (intersection r1 r)))))))))
      
      (define (subset? tree1 tree2)
        (or (empty? tree1)
            (and (<= (size tree1) (size tree2))
                 (with-n-node tree1
                              (lambda (k l r)
                                (cond ((elm< k (node-k tree2))
                                       (and (subset? l (node-l tree2))
                                            (member? k tree2)
                                            (subset? r tree2)))
                                      ((elm> k (node-k tree2))
                                       (and (subset? r (node-r tree2))
                                            (member? k tree2)
                                            (subset? l tree2)))
                                      (else
                                       (and (subset? l (node-l tree2))
                                            (subset? r (node-r tree2))))))))))
      (define (elements tree)
        (fold cons '() tree))
      
      (define (equal? tree1 tree2)
        (and (subset? tree1 tree2)
             (subset? tree2 tree1)))

      )))


